#pragma once

#include <arpa/inet.h>
#include <fcntl.h>
#include <netinet/tcp.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <cerrno>
#include <cstring>
#include <experimental/coroutine>
#include <optional>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <variant>
#include <vector>
#include "bio_ev.h"
#include "bio_util.h"

namespace bio {
using std::experimental::coroutine_handle;

template <typename T>
class async;

namespace priv {
template <typename T, bool terminating = false, typename A = async<T>>
struct promise_base {
   private:
    struct fs_awaitable {
        bool await_ready() const noexcept { return false; }
        template <typename P>
        coroutine_handle<> await_suspend(coroutine_handle<P> hdl) noexcept {
            return hdl.promise().awaiter;
        }
        void await_resume() noexcept {}
    };

   public:
    auto initial_suspend() noexcept { return std::experimental::suspend_always{}; }

    template <bool U = terminating, typename = std::enable_if_t<!U>>
    fs_awaitable final_suspend() noexcept {
        return fs_awaitable{};
    }

    template <bool U = terminating, typename = std::enable_if_t<U>>
    std::experimental::suspend_always final_suspend() noexcept {
        return std::experimental::suspend_always{};
    }

    template <bool U = terminating, typename = std::enable_if_t<!U>>
    void register_awaiter(coroutine_handle<> hdl) {
        awaiter = hdl;
    }

   private:
    // NON-OWNING handle to coroutine awaiting this one's promise
    coroutine_handle<> awaiter;
};

template <typename T, bool terminating = false, typename A = async<T>>
struct promise_type : promise_base<T, terminating> {
    auto get_return_object() noexcept {
        return A{coroutine_handle<promise_type>::from_promise(*this)};
    }
    void unhandled_exception() noexcept {
        result.template emplace<std::exception_ptr>(std::current_exception());
    }

    template <typename U = T, typename = std::enable_if_t<!std::is_void_v<U>>>
    void return_value(T&& value) {
        result.template emplace<T>(std::move(value));
    }
    template <typename U = T, typename = std::enable_if_t<!std::is_void_v<U>>>
    void return_value(T& value) {
        result.template emplace<T>(value);
    }

    template <typename U = T, typename = std::enable_if_t<!std::is_void_v<U>>>
    T& get_result() & {
        if (std::holds_alternative<std::exception_ptr>(result)) {
            std::rethrow_exception(std::get<std::exception_ptr>(result));
        } else if (std::holds_alternative<T>(result)) {
            return std::get<T>(result);
        } else {
            std::abort();
        }
    }

    template <typename U = T, typename = std::enable_if_t<!std::is_void_v<U>>>
    T&& get_result() && {
        if (std::holds_alternative<std::exception_ptr>(result)) {
            std::rethrow_exception(std::get<std::exception_ptr>(result));
        } else if (std::holds_alternative<T>(result)) {
            return std::move(std::get<T>(result));
        } else {
            std::abort();
        }
    }

   private:
    struct empty_t {};
    std::variant<empty_t, T, std::exception_ptr> result;
};

template <bool terminating, typename A>
struct promise_type<void, terminating, A> : promise_base<void, terminating, A> {
    auto get_return_object() noexcept {
        return A{coroutine_handle<promise_type>::from_promise(*this)};
    }

    void return_void() {}
    void unhandled_exception() noexcept {
        result.template emplace<std::exception_ptr>(std::current_exception());
    }

    void get_result() {
        if (std::holds_alternative<std::exception_ptr>(result)) {
            std::rethrow_exception(std::get<std::exception_ptr>(result));
        }
    }

   private:
    struct empty_t {};
    std::variant<empty_t, std::exception_ptr> result;
};

};  // namespace priv

template <typename T>
class [[nodiscard]] async {
   public:
    using promise_type = priv::promise_type<T>;
    using handle = coroutine_handle<promise_type>;

    async() noexcept : hdl(nullptr) {}
    explicit async(handle hdl) : hdl(hdl) {}

    async(async && other) noexcept : hdl(other.hdl) { other.hdl = nullptr; }

    async& operator=(async&& other) noexcept {
        if (&other == this) {
            return *this;
        }
        if (hdl) {
            hdl.destroy();
        }
        hdl = other.hdl;
        other.hdl = nullptr;
    }

    async(const async&) = delete;
    async& operator=(const async&) = delete;
    ~async() {
        if (hdl) {
            hdl.destroy();
        }
    }

    bool is_complete() { return !hdl || hdl.done(); }

    auto operator co_await() const& noexcept {
        struct temp : async_awaitable {
            decltype(auto) await_resume() {
                if (!this->hdl) {
                    std::abort();
                }
                return this->hdl.promise().get_result();
            }
        };
        return temp{hdl};
    }

    template <typename U = T, typename = std::enable_if_t<!std::is_void_v<U>>>
    auto operator co_await() const&& noexcept {
        struct temp : async_awaitable {
            decltype(auto) await_resume() {
                if (!this->hdl) {
                    std::abort();
                }
                return std::move(this->hdl.promise().get_result());
            }
        };
        return temp{hdl};
    }

   private:
    struct async_awaitable {
        async_awaitable(handle hdl) noexcept : hdl(hdl) {}

        bool await_ready() const noexcept { return !hdl || hdl.done(); }

        coroutine_handle<> await_suspend(coroutine_handle<> other) {
            hdl.promise().register_awaiter(other);
            return hdl;
        }

        handle hdl;
    };
    handle hdl;
};

namespace priv {

template <typename T>
class async_result {
   public:
    using promise_type = promise_type<T, true, async_result<T>>;
    using handle = coroutine_handle<promise_type>;

    async_result() noexcept : hdl(nullptr), ready(false) {}
    explicit async_result(handle hdl) : hdl(hdl), ready(true) {}

    async_result(async_result&& other) noexcept : hdl(other.hdl), ready(other.ready) {
        other.hdl = nullptr;
        other.ready = false;
    }

    async_result& operator=(async_result&& other) noexcept {
        if (&other == this) {
            return *this;
        }
        if (hdl) {
            hdl.destroy();
        }
        hdl = other.hdl;
        other.hdl = nullptr;
        ready = other.ready;
        other.ready = false;
        return *this;
    }

    async_result(const async_result&) = delete;
    async_result& operator=(const async_result&) = delete;
    ~async_result() {
        if (hdl) {
            hdl.destroy();
        }
    }

    bool is_ready() const noexcept { return hdl && !hdl.done() && ready; }
    bool is_done() const noexcept { return !hdl || hdl.done(); }
    void resume() {
        ready = false;
        hdl.resume();
    }

    promise_type& get_promise() { return hdl.promise(); }

   private:
    handle hdl;
    bool ready;
};

};  // namespace priv

thread_local struct ev_loop* loop = NULL;
thread_local struct ev_idle idler;
thread_local struct ev_check bg_runner;
thread_local std::vector<coroutine_handle<>> background_tasks;

namespace priv {
struct SignalHandler;
};

std::unordered_map<int, priv::SignalHandler> signals;

static void init() {
    if (loop == NULL) {
        // TODO: not multithread safe
        loop = EV_DEFAULT;
        ev_idle_init(&idler, [](struct ev_loop*, struct ev_idle*, int) {});
        ev_check_init(&bg_runner, [](struct ev_loop* loop, ev_check*, int) {
            log("Background task(s) is/are ready to run");
            ev_idle_stop(loop, &idler);
            ev_check_stop(loop, &bg_runner);
            std::vector<coroutine_handle<>> tmp = std::move(background_tasks);
            for (auto& t : tmp) {
                log("starting bg task");
                t.resume();
            }
            log("Finished invoking background tasks");
        });
    }
}

static inline struct ev_loop* get_loop() { return loop; }

static void destroy() {
    ev_idle_stop(loop, &idler);
    ev_check_stop(loop, &bg_runner);

    ev_loop_destroy(loop);
}

namespace priv {

struct SignalHandler {
    ev_signal w;
    std::function<void()> callback;
    SignalHandler(int which, std::function<void()> callback) : callback(callback) {
        ev_signal_init(&w, watcher_cb, which);
        w.data = this;
        ev_signal_start(get_loop(), &w);
    }
    ~SignalHandler() {
        ev_signal_stop(get_loop(), &w);
    }
    static void watcher_cb(struct ev_loop*, ev_signal* w, int) {
        SignalHandler* ptr = reinterpret_cast<SignalHandler*>(w->data);
        auto cb = ptr->callback;
        cb();
    }
};

};  // namespace priv

// TODO: ensure main thread only
// once this project gets support for multithreading / multiple event loops
void signal(int which, std::function<void()> callback) {
    if (signals.find(which) != signals.end()) {
        throw std::runtime_error{"Signal already registered"};
    }
    signals.emplace(std::piecewise_construct, std::forward_as_tuple(which),
                    std::forward_as_tuple(which, callback));
}

void signal_del(int which) {
    auto it = signals.find(which);
    if (it != signals.end()) {
        signals.erase(it);
    }
}

void set_nonblock(int fd) {
    int flags;
    if ((flags = fcntl(fd, F_GETFL)) < 0) {
        throw errno_exception{};
    }
    flags |= O_NONBLOCK;
    if (fcntl(fd, F_SETFL) < 0) {
        throw errno_exception{};
    }
}

auto read(int fd, size_t nbyte) {
    struct helper {
        int fd;
        bytes b;
        ssize_t result;
        error_t err;
        ev_io watcher;
        coroutine_handle<> to_resume;

        static void reader_cb(struct ev_loop* loop, ev_io* w, int) {
            helper* ptr = reinterpret_cast<helper*>(w->data);
            ptr->result = ::read(ptr->fd, ptr->b.data(), ptr->b.size());
            if (ptr->result < 0 && (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)) {
                // try again later
                return;
            } else if (ptr->result < 0) {
                ptr->err = errno;
            } else {
                ptr->b.resize(ptr->result);
            }
            ev_io_stop(loop, w);
            ptr->to_resume.resume();
        }
        bool await_ready() { return false; }
        void await_suspend(coroutine_handle<> hdl) {
            ev_io_init(&watcher, reader_cb, fd, EV_READ);
            to_resume = hdl;
            watcher.data = this;
            ev_io_start(get_loop(), &watcher);
        }
        bytes&& await_resume() {
            if (result < 0) {
                throw errno_exception(err);
            }
            return std::move(b);
        }
    };
    return helper{fd, {nbyte}, -1, 0, {}, nullptr};
}

class TcpClient {
   private:
    int fd;
    ev_io read_w;
    ev_io write_w;
    coroutine_handle<> read_waiter;
    coroutine_handle<> write_waiter;
    ssize_t result;
    error_t err;

    bytes* read_data;
    size_t read_offset;
    size_t read_len;

    static void read_cb(struct ev_loop* loop, ev_io* w, int) {
        TcpClient* ptr = reinterpret_cast<TcpClient*>(w->data);
        ptr->result = ::recv(ptr->fd, ptr->read_data->data() + ptr->read_offset, ptr->read_len, 0);
        if (ptr->result < 0 && (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)) {
            // try again later
            return;
        } else if (ptr->result < 0) {
            ptr->err = errno;
        }
        ev_io_stop(loop, w);
        ptr->read_waiter.resume();
    }

    static void write_cb(struct ev_loop* loop, ev_io* w, int) {
        TcpClient* ptr = reinterpret_cast<TcpClient*>(w->data);
        ev_io_stop(loop, w);
        ptr->write_waiter.resume();
    }

    void do_move(TcpClient& other) {
        log("Moving TcpClient");
        fd = other.fd;
        other.fd = -1;
        read_waiter = std::move(other.read_waiter);
        other.read_waiter = nullptr;
        write_waiter = std::move(other.write_waiter);
        other.write_waiter = nullptr;

        ev_io_stop(get_loop(), &other.read_w);
        ev_io_stop(get_loop(), &other.write_w);

        ev_io_init(&read_w, read_cb, fd, EV_READ);
        read_w.data = this;
        ev_io_init(&write_w, read_cb, fd, EV_WRITE);
        write_w.data = this;
        if (read_waiter) {
            ev_io_start(get_loop(), &read_w);
        }
        if (write_waiter) {
            ev_io_start(get_loop(), &write_w);
        }
    }

   public:
    TcpClient(int fd) : fd(fd) {
        log("Initializing TcpClient");
        try {
            set_nonblock(fd);
        } catch (errno_exception const& e) {
            log("ran into errno on nonblock");
            log(e.what());
        }
        ev_io_init(&read_w, read_cb, fd, EV_READ);
        read_w.data = this;
        ev_io_init(&write_w, read_cb, fd, EV_WRITE);
        write_w.data = this;
        read_waiter = nullptr;
        write_waiter = nullptr;
        log("Done initializing TcpClient");
    }

    TcpClient(const TcpClient&) = delete;
    TcpClient& operator=(const TcpClient&) = delete;

    TcpClient(TcpClient&& other) { do_move(other); }
    TcpClient& operator=(TcpClient&& other) {
        do_move(other);
        return *this;
    }

    auto read(size_t nbyte) {
        struct helper {
            TcpClient& client;
            bytes data;
            bool await_ready() { return false; }

            void await_suspend(coroutine_handle<> hdl) {
                client.read_waiter = hdl;
                client.read_data = &data;
                client.read_offset = 0;
                client.read_len = data.size();
                ev_io_start(get_loop(), &client.read_w);
            }

            bytes await_resume() {
                if (client.result < 0) {
                    throw errno_exception(client.err);
                } else {
                    data.resize(client.result);
                }
                return std::move(data);
            }
        };
        return helper{*this, {nbyte}};
    }

    auto read(bytes& bytes, size_t offset, size_t len) {
        assert(offset + len <= bytes.size());

        struct helper {
            TcpClient& client;
            class bytes& data;
            size_t offset;
            size_t len;
            bool await_ready() { return false; }
            void await_suspend(coroutine_handle<> hdl) {
                client.read_waiter = hdl;
                client.read_data = &data;
                client.read_offset = offset;
                client.read_len = len;
                ev_io_start(get_loop(), &client.read_w);
            }
            size_t await_resume() {
                if (client.result < 0) {
                    throw errno_exception{client.err};
                }
                return client.result;
            }
        };
        return helper{*this, bytes, offset, len};
    }

    auto write_ready() {
        struct helper {
            TcpClient& client;
            bool await_ready() { return false; }
            void await_suspend(coroutine_handle<> hdl) {
                client.write_waiter = hdl;
                ev_io_start(get_loop(), &client.write_w);
            }
            void await_resume() {}
        };
        return helper{*this};
    }

    auto write(bytes b) { return write(b, 0, b.size()); }

    // CLANG BUG: workaround for clang bug currently still present in trunk
    // Coroutine + member variable + ubsan = boom :(
    __attribute__((no_sanitize("undefined"))) async<void> write(bytes b, size_t offset,
                                                                size_t len) {
        assert(offset + len <= b.size());

        if (len == 0) {
            co_return;
        }
        size_t ptr = offset;

        while (ptr < offset + len) {
            ssize_t res = ::send(fd, b.data() + ptr, offset + len - ptr, 0);
            if (res < 0 && (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)) {
                co_await write_ready();
            } else if (res < 0) {
                throw errno_exception{};
            }
            ptr += res;
        }

        co_return;
    }

    void shutdown() {
        if (fd > -1) {
            log("Shutting down TcpClient");
            // TODO: errors
            ::shutdown(fd, SHUT_RDWR);
            close(fd);
            fd = -1;
            ev_io_stop(get_loop(), &read_w);
        }
    }

    ~TcpClient() { shutdown(); }
};

class TcpServer {
    using connect_cb = std::function<void(TcpClient, struct sockaddr_in6, TcpServer&)>;

   private:
    static constexpr const int MAX_PENDING = 8;
    int fd;
    ev_io accept_w;
    connect_cb cb;
    coroutine_handle<> waiter;

    static void accept_cb(struct ev_loop*, ev_io* w, int) {
        log("Accept callback called");
        TcpServer* serv = reinterpret_cast<TcpServer*>(w->data);
        struct sockaddr_in6 client_addr;
        socklen_t addrlen = sizeof(client_addr);
        int res = accept(serv->fd, reinterpret_cast<struct sockaddr*>(&client_addr), &addrlen);
        if (res < 0 && (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)) {
            // try again later
            return;
        } else if (res < 0) {
            // TODO: handle accept error
            log("accept error");
            log(strerror(errno));
            return;
        }

        log("Client accepted, Invoking callback");
        serv->cb(TcpClient{res}, client_addr, *serv);
    }

    void do_move(TcpServer& other) {
        fd = other.fd;
        other.fd = -1;
        cb = std::move(other.cb);
        waiter = std::move(other.waiter);
        ev_io_stop(get_loop(), &other.accept_w);

        ev_io_init(&accept_w, accept_cb, fd, EV_READ);
        accept_w.data = this;
        ev_io_start(get_loop(), &accept_w);
    }

   public:
    TcpServer(uint16_t port, connect_cb cb) : cb(cb) {
        log("Tcp server init");
        if ((fd = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP)) < 0) {
            throw errno_exception{};
        }
        struct sockaddr_in6 addr = {};
        addr.sin6_family = AF_INET6;
        addr.sin6_addr = IN6ADDR_ANY_INIT;
        addr.sin6_port = htons(port);

        set_nonblock(fd);

        int arg = 1;
        if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &arg, sizeof(int)) < 0) {
            throw errno_exception{};
        }

        if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &arg, sizeof(int)) < 0) {
            throw errno_exception{};
        }

        int no = 0;
        if (setsockopt(fd, IPPROTO_IPV6, IPV6_V6ONLY, &no, sizeof(no))) {
            throw errno_exception{};
        }

        if (bind(fd, reinterpret_cast<struct sockaddr*>(&addr), sizeof(addr)) < 0) {
            throw errno_exception{};
        }

        if (listen(fd, MAX_PENDING) < 0) {
            throw errno_exception{};
        }
        log("Socket set up");

        ev_io_init(&accept_w, accept_cb, fd, EV_READ);
        accept_w.data = this;
        ev_io_start(get_loop(), &accept_w);
        log("Accept watcher started");
    }

    TcpServer(const TcpServer& other) = delete;
    TcpServer& operator=(const TcpServer& other) = delete;

    TcpServer(TcpServer&& other) {
        log("Moving TcpServer");
        do_move(other);
    }

    TcpServer& operator=(TcpServer&& other) {
        log("Moving TcpServer");
        do_move(other);
        return *this;
    }

    void shutdown() {
        if (fd > -1) {
            log("Shutting down TcpServer");
            // TODO: errors?
            close(fd);
            ev_io_stop(get_loop(), &accept_w);
            fd = -1;
            if (waiter) {
                waiter.resume();
            }
        }
    }

    auto wait_shutdown() {
        struct helper {
            TcpServer& serv;
            bool await_ready() { return serv.fd <= -1; }
            void await_suspend(coroutine_handle<> hdl) { serv.waiter = hdl; }
            void await_resume() {}
        };
        return helper{*this};
    }

    ~TcpServer() { shutdown(); }
};

auto yield() {
    struct yield_awaitable {
        bool await_ready() { return false; }
        void await_suspend(coroutine_handle<> hdl) {
            background_tasks.push_back(hdl);
            ev_idle_start(get_loop(), &idler);
            ev_check_start(get_loop(), &bg_runner);
        }
        void await_resume() {}
    };
    return yield_awaitable{};
}

template <typename T>
void background(async<T> as) {
    log("Backgrounding task");
    // A mini-coroutine-context system that lives on the heap and self destructs after exiting
    struct background_task {
        struct promise_type;
        using handle = coroutine_handle<promise_type>;
        struct promise_type {
            auto get_return_object() { return background_task{handle::from_promise(*this)}; }
            auto initial_suspend() { return std::experimental::suspend_always{}; }
            auto final_suspend() {
                struct awaitable {
                    bool await_ready() { return false; }
                    void await_suspend(coroutine_handle<> hdl) { hdl.destroy(); }
                    void await_resume() {}
                };
                return awaitable{};
            }
            void return_void() {}
            void unhandled_exception() { log("Unhandled exception in background task"); }
        };
        handle hdl;
    };
    auto wrapper = [](async<T> as) -> background_task {
        log("background task: awaiting wrapped task");
        co_await as;
        log("background task: done awaiting");
        co_return;
    };
    background_tasks.push_back(wrapper(std::move(as)).hdl);
    ev_idle_start(get_loop(), &idler);
    ev_check_start(get_loop(), &bg_runner);
}

template <typename T>
T run(async<T>& as) {
    struct helper {
        priv::async_result<T> wrapper(async<T>& as) { co_return co_await as; }
    };
    helper h;
    priv::async_result<T> res = h.wrapper(as);

    ev_prepare prp;
    ev_prepare_init(&prp, [](struct ev_loop* loop, ev_prepare* w, int) {
        priv::async_result<T>* ptr = reinterpret_cast<priv::async_result<T>*>(w->data);
        ev_prepare_stop(loop, w);
        ptr->resume();
    });
    prp.data = &res;
    ev_prepare_start(get_loop(), &prp);

    ev_run(get_loop(), 0);

    signals.clear();
    background_tasks.clear();

    return res.get_promise().get_result();
}

template <typename T>
T run(async<T>&& as) {
    return run(as);
}

}  // namespace bio
